package middleware

import (
	"context"
	"fmt"
	"github.com/mark3labs/mcp-go/server"
	"go.uber.org/zap"
	"mcp-demo/internal/global"
	"mcp-demo/internal/model/vo/common/result"
	"mcp-demo/internal/utils/jsonutil"
	"runtime/debug"
	"time"

	"github.com/mark3labs/mcp-go/mcp"
)

func ToolErrorMiddleware(next server.ToolHandlerFunc) server.ToolHandlerFunc {
	return func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		// 1. 超时控制
		timeoutCtx, cancel := context.WithTimeout(ctx, 15*time.Second)
		defer cancel()

		// 2. 使用通道处理panic和正常结果
		resultChan := make(chan *mcp.CallToolResult, 1)
		errorChan := make(chan error, 1)

		go func() {
			defer func() {
				if r := recover(); r != nil {
					global.L.Error(fmt.Sprintf("[MCP Recovery from panic])\n%v%s", r, debug.Stack()),
						zap.Any("error", r),
						zap.String("stack", string(debug.Stack())),
					)
					jsonString, _ := jsonutil.GetJsonString(result.FailResult())
					errorChan <- nil // 阻止panic继续传播
					resultChan <- &mcp.CallToolResult{
						Content: []mcp.Content{
							mcp.TextContent{
								Type: "text",
								Text: jsonString,
							},
						},
					}
				}
			}()

			// 实际调用工具
			res, err := next(timeoutCtx, req)
			if err != nil {
				jsonString, _ := jsonutil.GetJsonString(result.FailResult())
				errorChan <- nil // 阻止error继续传播
				resultChan <- &mcp.CallToolResult{
					Content: []mcp.Content{
						mcp.TextContent{
							Type: "text",
							Text: jsonString,
						},
					},
				}
				return
			}
			errorChan <- nil
			resultChan <- res
		}()

		// 等待结果或超时
		select {
		case <-timeoutCtx.Done():
			jsonString, _ := jsonutil.GetJsonString(result.FailResult())
			return &mcp.CallToolResult{
				Content: []mcp.Content{
					mcp.TextContent{
						Type: "text",
						Text: jsonString,
					},
				},
			}, nil
		case err := <-errorChan:
			return <-resultChan, err
		case res := <-resultChan:
			return res, <-errorChan
		}
	}
}
